#include "stdafx.h"
#include "RegisterClass.h"

#define GUID_CCH	39	// Characters in string form of guid, including '\0'

CRegisterClass::CRegisterClass(REFCLSID clsid, LPCTSTR pszTypeName, HKEY hkeyRoot)
	: m_hkeyRoot(hkeyRoot)
{
	// The constructor calculates three predefined symbols:
	// ClassID  %1 = "{4d3450049-1e19-..."
	// TypeName %2 = "MyComponent.Object"
	// FullPath %3 = "c:\mycode\debug\myobject.dll"

	// Format class ID as a string
	OLECHAR szClassID[GUID_CCH];
	int cchGuid = ::StringFromGUID2(clsid, szClassID, GUID_CCH);

	// Convert to multibyte from unicode
	CString strClassID(szClassID);

	ASSERT(cchGuid == GUID_CCH);	// Did StringFromGUID2 work?
	if (cchGuid != GUID_CCH)
	{
		TRACE("StringFromGUID2 failed\n");
		return;
	}

	CString strPathName;
	TCHAR szLongPathName[_MAX_PATH];
	::GetModuleFileName(AfxGetResourceHandle(), szLongPathName, _MAX_PATH);
	::GetShortPathName(szLongPathName,
		strPathName.GetBuffer(_MAX_PATH), _MAX_PATH);
	strPathName.ReleaseBuffer();

	CString strTypeName(pszTypeName);

	m_symbols.Add(strClassID);
	m_symbols.Add(strTypeName);
	m_symbols.Add(strPathName);
}

LONG AFXAPI _RecursiveRegDeleteKey(HKEY hParentKey, LPCTSTR szKeyName);

BOOL CRegisterClass::RemoveKey(LPCTSTR key)
{
	CString strKey;

	// Build an array of pointers to the symbols
	int nSymbols = m_symbols.GetSize();
	CString* pStrings = m_symbols.GetData();

	LPCTSTR* ppSymbols = new LPCTSTR[nSymbols];
	for (int i=0; i < nSymbols; i++)
		ppSymbols[i] = (LPCTSTR) pStrings[i];

	AfxFormatStrings(strKey, key, ppSymbols, nSymbols);
	BOOL bRet = _RecursiveRegDeleteKey(m_hkeyRoot, strKey) == ERROR_SUCCESS;

	delete [] ppSymbols;
	return bRet;
}

void CRegisterClass::AddEntry(LPCTSTR key, LPCTSTR value, LPCTSTR valueName)
{
	ASSERT(AfxIsValidString(key));
	ASSERT(AfxIsValidString(value));
	ASSERT(valueName == NULL || AfxIsValidString(valueName));

	m_keys.Add(key);
	m_values.Add(value);
	m_valueNames.Add(valueName);
}

int CRegisterClass::AddSymbol(LPCTSTR pSymbol)
{
	return m_symbols.Add(pSymbol);
}

int CRegisterClass::AddSymbolAlias(LPCTSTR pSymbol, LPCTSTR pAlias)
{
	TCHAR szValue[MAX_PATH];
	LONG lSize = sizeof(szValue);

	// Don't replace pSymbol's existing value
	if (RegQueryValue(HKEY_CLASSES_ROOT, pSymbol, szValue, &lSize) == ERROR_SUCCESS)
	{
		// If alias already exists for symbol, use it instead of requested alias (pAlias)
		if (lSize > 1)
			return AddSymbol(szValue);
	}
	
	// SetRegKey will add pSymbol if it doesn't exist with pAlias as its value
	SetRegKey(pSymbol, pAlias);
	return AddSymbol(pAlias);
}

BOOL CRegisterClass::Register(BOOL bRegister)
{
	// Build an array of pointers to the symbols
	int nSymbols = m_symbols.GetSize();
	CString* pStrings = m_symbols.GetData();

	LPCTSTR* ppSymbols = new LPCTSTR[nSymbols];
	for (int i=0; i < nSymbols; i++)
		ppSymbols[i] = (LPCTSTR) pStrings[i];

	BOOL bRet = Helper(bRegister, ppSymbols, nSymbols);
	delete ppSymbols;		// not delete [] pSymbols
	return bRet;
}


BOOL CRegisterClass::Helper(BOOL bRegister, LPCTSTR const* ppSymbols, int nSymbols)
{
	ASSERT(nSymbols != 0);

	CString strKey;
	CString strValue;
	CString strValueName;	// (can be NULL)

	BOOL bResult = TRUE;
	int nItems = m_keys.GetSize();

	for (int i = 0; i < nItems; i++)
	{
		AfxFormatStrings(strKey, m_keys[i], ppSymbols, nSymbols);
		AfxFormatStrings(strValue, m_values[i], ppSymbols, nSymbols);
		AfxFormatStrings(strValueName, m_valueNames[i], ppSymbols, nSymbols);

		if (strKey.IsEmpty())
			continue;

		if (bRegister)
		{
			if (!SetRegKey(strKey, strValue, strValueName))
			{
				TRACE("Error: failed to set key '%s' to value '%s'.\n",
					(LPCTSTR) strKey, (LPCTSTR) strValue);
				bResult = FALSE;
				break;
			}
		} else {
			if (!DelRegKey(strKey, strValueName))
			{
				TRACE("Failed to delete key %s\n", (LPCSTR) strKey);
				bResult = FALSE;
				break;
			}
		}
	}

	return bResult;
}

BOOL CRegisterClass::DelRegKey(LPCTSTR lpszKey, LPCTSTR lpszValueName)
{
	if (lpszValueName == NULL || lstrlen(lpszValueName) == 0)
	{
		ASSERT(AfxIsValidString(lpszKey));
		return (::RegDeleteKey(m_hkeyRoot, lpszKey) == ERROR_SUCCESS);
	} else {
		HKEY hKey;
		if (::RegOpenKey(m_hkeyRoot, lpszKey, &hKey) == ERROR_SUCCESS)
		{
			LONG lResult = ::RegDeleteValue(hKey, lpszValueName);
			if(::RegCloseKey(hKey) == ERROR_SUCCESS && lResult == ERROR_SUCCESS)
				return TRUE;
		}
		return FALSE;
	}
}

BOOL CRegisterClass::SetRegKey(LPCTSTR lpszKey, LPCTSTR lpszValue, LPCTSTR lpszValueName)
{
	if (lpszValueName == NULL || lstrlen(lpszValueName) == 0)
	{
		if (::RegSetValue(m_hkeyRoot, lpszKey, REG_SZ,
			  lpszValue, lstrlen(lpszValue)) != ERROR_SUCCESS)
		{
			TRACE1("Warning: registration database update failed for key '%s'.\n",
				lpszKey);
			return FALSE;
		}
		return TRUE;
	}
	else
	{
		HKEY hKey;

		if(::RegCreateKey(m_hkeyRoot, lpszKey, &hKey) == ERROR_SUCCESS)
		{
			LONG lResult = ::RegSetValueEx(hKey, lpszValueName, 0, REG_SZ,
				(CONST BYTE*)lpszValue, lstrlen(lpszValue) + sizeof(TCHAR) );

			if(::RegCloseKey(hKey) == ERROR_SUCCESS && lResult == ERROR_SUCCESS)
				return TRUE;
		}
		TRACE1("Warning: registration database update failed for key '%s'.\n", lpszKey);
		return FALSE;
	}
}

// Under Win32, a reg key may not be deleted unless it is empty.
// Thus, to delete a tree,  one must recursively enumerate and
// delete all of the sub-keys.

LONG AFXAPI _RecursiveRegDeleteKey(HKEY hParentKey, LPCTSTR szKeyName)
{
	DWORD   dwIndex = 0L;
	TCHAR   szSubKeyName[256];
	HKEY    hCurrentKey;
	DWORD   dwResult;

	if ((dwResult = RegOpenKey(hParentKey, szKeyName, &hCurrentKey)) ==
		ERROR_SUCCESS)
	{
		// Remove all subkeys of the key to delete
		while ((dwResult = RegEnumKey(hCurrentKey, 0, szSubKeyName, 255)) ==
			ERROR_SUCCESS)
		{
			if ((dwResult = _RecursiveRegDeleteKey(hCurrentKey,
				szSubKeyName)) != ERROR_SUCCESS)
				break;
		}

		// If all went well, we should now be able to delete the requested key
		if ((dwResult == ERROR_NO_MORE_ITEMS) || (dwResult == ERROR_BADKEY))
		{
			dwResult = RegDeleteKey(hParentKey, szKeyName);
		}
	}

	RegCloseKey(hCurrentKey);
	return dwResult;
}

